import torch

from time import time
from torch.utils.data import DataLoader
from config import Config
from util.model import BERT
from util.dataTool import SeqDataset
from util.trainer import train, evaluate

if __name__ == '__main__':
    config = Config()

    tik = time()

    print("Data loading...")
    train_data = SeqDataset(config.train_path, config)
    dev_data = SeqDataset(config.dev_path, config)

    train_loader = DataLoader(train_data, batch_size=config.batch_size,
                              shuffle=True, pin_memory=True, num_workers=4, drop_last=True)
    dev_loader = DataLoader(dev_data, batch_size=config.batch_size,
                            shuffle=True, pin_memory=True, num_workers=4, drop_last=False)

    print("Model loading...")
    model = BERT(config).to(config.device)

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(),
                                  lr=config.learning_rate, weight_decay=config.weight_decay)

    print("Loading time:", time() - tik)
    tik = time()

    print("Training...")
    model = train(model,
                  loader=train_loader,
                  criterion=criterion,
                  optimizer=optimizer,
                  config=config)

    print("Training time:", time() - tik)
    tik = time()

    print("Testing...", round(len(dev_data) / config.batch_size))
    evaluate(model, dev_loader, config)
    print("Testing time:", time() - tik)